package com.liming.common.util;

import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author liming
 * @date 2025/1/10
 */
@Slf4j
public class ThreadLocalUtils {

    // 使用 ConcurrentHashMap 存储 ThreadLocal 变量
    private static final Map<String, ThreadLocal<Object>> THREAD_LOCAL_MAP = new ConcurrentHashMap<>();

    private ThreadLocalUtils() {
        throw new UnsupportedOperationException("Utility class should not be instantiated");
    }

    /**
     * 设置 ThreadLocal 变量
     *
     * @param key   变量的键
     * @param value 变量的值
     */
    public static void set(String key, Object value) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        ThreadLocal<Object> threadLocal = THREAD_LOCAL_MAP.computeIfAbsent(key, k -> new ThreadLocal<>());
        threadLocal.set(value);
    }

    /**
     * 获取 ThreadLocal 变量
     *
     * @param key 变量的键
     * @return 变量的值
     */
    @SuppressWarnings("unchecked")
    public static <T> T get(String key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        ThreadLocal<Object> threadLocal = THREAD_LOCAL_MAP.get(key);
        if (threadLocal == null) {
            log.warn("ThreadLocal with key {} not found", key);
            return null;
        }
        return (T) threadLocal.get();
    }

    /**
     * 获取 ThreadLocal 变量，并提供默认值
     *
     * @param key          变量的键
     * @param defaultValue 默认值
     * @return 变量的值或默认值
     */
    @SuppressWarnings("unchecked")
    public static <T> T getOrDefault(String key, T defaultValue) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        ThreadLocal<Object> threadLocal = THREAD_LOCAL_MAP.get(key);
        if (threadLocal == null) {
            log.warn("ThreadLocal with key {} not found, returning default value", key);
            return defaultValue;
        }
        Object value = threadLocal.get();
        return value != null ? (T) value : defaultValue;
    }

    /**
     * 移除 ThreadLocal 变量
     *
     * @param key 变量的键
     */
    public static void remove(String key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        ThreadLocal<Object> threadLocal = THREAD_LOCAL_MAP.get(key);
        if (threadLocal != null) {
            threadLocal.remove();
        } else {
            log.warn("ThreadLocal with key {} not found", key);
        }
    }

    /**
     * 清除所有 ThreadLocal 变量
     */
    public static void clearAll() {
        for (ThreadLocal<Object> threadLocal : THREAD_LOCAL_MAP.values()) {
            threadLocal.remove();
        }
        THREAD_LOCAL_MAP.clear();
    }
}
